from model.Ark import ArkModel

from langchain.agents import create_agent
from langchain.agents.middleware import before_model, before_agent, AgentState, PIIMiddleware
from langgraph.runtime import Runtime
from langchain.tools import tool

class CustomerAgentState(AgentState):
    user_id: str

@tool
def get_weather(city: str):
    """ 获取一个城市的天气情况
    Args:
        city: 城市名称
    """
    print(city)
    return '晴天'

""" 可以动态的使用 jump_to 来决定是否需要继续往下执行流程。
    
    jump_to: 有三个有效值：model、tool、end 
"""
@before_agent(can_jump_to=['end'])
def before_agent_handler(state: CustomerAgentState, runtime: Runtime):
    print('---> before agent')
    if state['user_id'] == '123':
        return None
    else:
        return {
            'jump_to': 'end',
            'messages': [
                { 'role': 'ai', 'content': '没有权限访问!' }
            ]
        }

@before_model
def before_model_handler(state: CustomerAgentState, runtime: Runtime):
    print('---> before model')
    return None

def run():
    agent = create_agent(
        model=ArkModel().model,
        middleware=[
            before_agent_handler,
            before_model_handler,
            ## 当返回的消息内容包含类型为："email", "credit_card", "ip", "mac_address", "url"时，可对其进行编码或报错的处理。
            PIIMiddleware(
                'url',
                strategy='mask'
            )
        ],
        tools=[get_weather],
        state_schema=CustomerAgentState
    )

    response = agent.invoke({ 'messages': [{ 'role': 'user', 'content': '深圳得今天的天气怎么样?, 链接地址为: http://127.0.0.1:3000' }], 'user_id': '123' })
    print(response['messages'][-1].content)
    print(response)